/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.fnexecution.environment;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnExternalWorkerPoolGrpc;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.Environment;

import com.bff.gaia.unified.runners.core.construction.UnifiedUrns;

import com.bff.gaia.unified.runners.fnexecution.GrpcFnServer;

import com.bff.gaia.unified.runners.fnexecution.artifact.ArtifactRetrievalService;

import com.bff.gaia.unified.runners.fnexecution.control.ControlClientPool;

import com.bff.gaia.unified.runners.fnexecution.control.FnApiControlClientPoolService;

import com.bff.gaia.unified.runners.fnexecution.control.InstructionRequestHandler;

import com.bff.gaia.unified.runners.fnexecution.logging.GrpcLoggingService;

import com.bff.gaia.unified.runners.fnexecution.provisioning.StaticGrpcProvisionService;

import com.bff.gaia.unified.sdk.fn.IdGenerator;

import com.bff.gaia.unified.sdk.fn.channel.ManagedChannelFactory;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.time.Duration;

import java.util.concurrent.TimeoutException;



/** An {@link EnvironmentFactory} which requests workers via the given URL in the Environment. */

public class ExternalEnvironmentFactory implements EnvironmentFactory {



  private static final Logger LOG = LoggerFactory.getLogger(ExternalEnvironmentFactory.class);



  public static ExternalEnvironmentFactory create(

      GrpcFnServer<FnApiControlClientPoolService> controlServiceServer,

      GrpcFnServer<GrpcLoggingService> loggingServiceServer,

      GrpcFnServer<ArtifactRetrievalService> retrievalServiceServer,

      GrpcFnServer<StaticGrpcProvisionService> provisioningServiceServer,

      ControlClientPool.Source clientSource,

      IdGenerator idGenerator) {

    return new ExternalEnvironmentFactory(

        controlServiceServer,

        loggingServiceServer,

        retrievalServiceServer,

        provisioningServiceServer,

        idGenerator,

        clientSource);

  }



  private final GrpcFnServer<FnApiControlClientPoolService> controlServiceServer;

  private final GrpcFnServer<GrpcLoggingService> loggingServiceServer;

  private final GrpcFnServer<ArtifactRetrievalService> retrievalServiceServer;

  private final GrpcFnServer<StaticGrpcProvisionService> provisioningServiceServer;

  private final IdGenerator idGenerator;

  private final ControlClientPool.Source clientSource;



  private ExternalEnvironmentFactory(

      GrpcFnServer<FnApiControlClientPoolService> controlServiceServer,

      GrpcFnServer<GrpcLoggingService> loggingServiceServer,

      GrpcFnServer<ArtifactRetrievalService> retrievalServiceServer,

      GrpcFnServer<StaticGrpcProvisionService> provisioningServiceServer,

      IdGenerator idGenerator,

      ControlClientPool.Source clientSource) {

    this.controlServiceServer = controlServiceServer;

    this.loggingServiceServer = loggingServiceServer;

    this.retrievalServiceServer = retrievalServiceServer;

    this.provisioningServiceServer = provisioningServiceServer;

    this.idGenerator = idGenerator;

    this.clientSource = clientSource;

  }



  /** Creates a new, active {@link RemoteEnvironment} backed by an unmanaged worker. */

  @Override

  public RemoteEnvironment createEnvironment(Environment environment) throws Exception {

    Preconditions.checkState(

        environment

            .getUrn()

            .equals(UnifiedUrns.getUrn(RunnerApi.StandardEnvironments.Environments.EXTERNAL)),

        "The passed environment does not contain an ExternalPayload.");

    final RunnerApi.ExternalPayload externalPayload =

        RunnerApi.ExternalPayload.parseFrom(environment.getPayload());

    final String workerId = idGenerator.getId();



    UnifiedFnApi.NotifyRunnerAvailableRequest notifyRunnerAvailableRequest =

        UnifiedFnApi.NotifyRunnerAvailableRequest.newBuilder()

            .setWorkerId(workerId)

            .setControlEndpoint(controlServiceServer.getApiServiceDescriptor())

            .setLoggingEndpoint(loggingServiceServer.getApiServiceDescriptor())

            .setArtifactEndpoint(retrievalServiceServer.getApiServiceDescriptor())

            .setProvisionEndpoint(provisioningServiceServer.getApiServiceDescriptor())

            .putAllParams(externalPayload.getParamsMap())

            .build();



    LOG.debug("Requesting worker ID {}", workerId);

    UnifiedFnApi.NotifyRunnerAvailableResponse notifyRunnerAvailableResponse =

        UnifiedFnExternalWorkerPoolGrpc.newBlockingStub(

                ManagedChannelFactory.createDefault().forDescriptor(externalPayload.getEndpoint()))

            .notifyRunnerAvailable(notifyRunnerAvailableRequest);

    if (!notifyRunnerAvailableResponse.getError().isEmpty()) {

      throw new RuntimeException(notifyRunnerAvailableResponse.getError());

    }



    // Wait on a client from the gRPC server.

    InstructionRequestHandler instructionHandler = null;

    while (instructionHandler == null) {

      try {

        instructionHandler = clientSource.take(workerId, Duration.ofMinutes(2));

      } catch (TimeoutException timeoutEx) {

        LOG.info(

            "Still waiting for startup of environment from {} for worker id {}",

            externalPayload.getEndpoint().getUrl(),

            workerId);

      } catch (InterruptedException interruptEx) {

        Thread.currentThread().interrupt();

        throw new RuntimeException(interruptEx);

      }

    }

    final InstructionRequestHandler finalInstructionHandler = instructionHandler;



    return new RemoteEnvironment() {

      @Override

      public Environment getEnvironment() {

        return environment;

      }



      @Override

      public InstructionRequestHandler getInstructionRequestHandler() {

        return finalInstructionHandler;

      }

    };

  }



  /** Provider of ExternalEnvironmentFactory. */

  public static class Provider implements EnvironmentFactory.Provider {

    @Override

    public EnvironmentFactory createEnvironmentFactory(

        GrpcFnServer<FnApiControlClientPoolService> controlServiceServer,

        GrpcFnServer<GrpcLoggingService> loggingServiceServer,

        GrpcFnServer<ArtifactRetrievalService> retrievalServiceServer,

        GrpcFnServer<StaticGrpcProvisionService> provisioningServiceServer,

        ControlClientPool clientPool,

        IdGenerator idGenerator) {

      return create(

          controlServiceServer,

          loggingServiceServer,

          retrievalServiceServer,

          provisioningServiceServer,

          clientPool.getSource(),

          idGenerator);

    }

  }

}